use std::cmp::Ordering;
use std::mem;

/// Type of heap: min-heap or max-heap.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BinaryHeapType {
    Min,
    Max,
}

/// A binary heap data structure.
pub struct BinaryHeap<T> {
    heap_type: BinaryHeapType,
    values: Vec<T>,
    compare_func: fn(&T, &T) -> Ordering,
}

impl<T> BinaryHeap<T> {
    /// Creates a new binary heap with the specified type and comparison function.
    ///
    /// # Arguments
    ///
    /// * `heap_type` - The type of heap (min-heap or max-heap).
    /// * `compare_func` - A function to compare two values.
    ///
    /// # Returns
    ///
    /// A new binary heap.
    pub fn new(heap_type: BinaryHeapType, compare_func: fn(&T, &T) -> Ordering) -> Self {
        BinaryHeap {
            heap_type,
            values: Vec::with_capacity(16),
            compare_func,
        }
    }

    /// Inserts a value into the heap.
    ///
    /// # Arguments
    ///
    /// * `value` - The value to insert.
    ///
    /// # Returns
    ///
    /// `true` if the insertion was successful.
    pub fn insert(&mut self, value: T) -> bool {
        let mut index = self.values.len();
        self.values.push(value);

        // Percolate the value up to the top of the heap.
        while index > 0 {
            let parent = (index - 1) / 2;
            if self.compare(&self.values[parent], &self.values[index]) == Ordering::Less {
                break;
            }
            self.values.swap(index, parent);
            index = parent;
        }

        true
    }

    /// Removes and returns the top value from the heap.
    ///
    /// # Returns
    ///
    /// The top value, or `None` if the heap is empty.
    pub fn pop(&mut self) -> Option<T> {
        if self.values.is_empty() {
            return None;
        }

        let result = self.values.swap_remove(0);
        let mut index = 0;

        // Percolate the new top value down.
        while index < self.values.len() {
            let child1 = index * 2 + 1;
            let child2 = index * 2 + 2;
            let mut next_index = index;

            if child1 < self.values.len()
                && self.compare(&self.values[child1], &self.values[index]) == Ordering::Less
            {
                next_index = child1;
            }
            if child2 < self.values.len()
                && self.compare(&self.values[child2], &self.values[next_index]) == Ordering::Less
            {
                next_index = child2;
            }

            if next_index == index {
                break;
            }

            self.values.swap(index, next_index);
            index = next_index;
        }

        Some(result)
    }

    /// Returns the number of entries in the heap.
    ///
    /// # Returns
    ///
    /// The number of entries.
    pub fn num_entries(&self) -> usize {
        self.values.len()
    }

    /// Compares two values based on the heap type.
    fn compare(&self, data1: &T, data2: &T) -> Ordering {
        let result = (self.compare_func)(data1, data2);
        match self.heap_type {
            BinaryHeapType::Min => result,
            BinaryHeapType::Max => result.reverse(),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_min_heap() {
        let mut heap = BinaryHeap::new(BinaryHeapType::Min, |a: &i32, b: &i32| a.cmp(b));
        heap.insert(3);
        heap.insert(1);
        heap.insert(4);
        heap.insert(1);
        heap.insert(5);

        assert_eq!(heap.pop(), Some(1));
        assert_eq!(heap.pop(), Some(1));
        assert_eq!(heap.pop(), Some(3));
        assert_eq!(heap.pop(), Some(4));
        assert_eq!(heap.pop(), Some(5));
        assert_eq!(heap.pop(), None);
    }

    #[test]
    fn test_max_heap() {
        let mut heap = BinaryHeap::new(BinaryHeapType::Max, |a: &i32, b: &i32| a.cmp(b));
        heap.insert(3);
        heap.insert(1);
        heap.insert(4);
        heap.insert(1);
        heap.insert(5);

        assert_eq!(heap.pop(), Some(5));
        assert_eq!(heap.pop(), Some(4));
        assert_eq!(heap.pop(), Some(3));
        assert_eq!(heap.pop(), Some(1));
        assert_eq!(heap.pop(), Some(1));
        assert_eq!(heap.pop(), None);
    }
}
